from faker import Faker
import random
from datetime import datetime, timedelta

fake = Faker('zh_CN')

risk_levels = ["高危", "低危", "无法评估", "无需评估"]
patient_types = ["门诊", "住院"]
sex_list = ["男", "女"]


def generate_patient(i):
    age = random.randint(18, 90)
    sex = random.choice(sex_list)
    birthday = (datetime.now() - timedelta(days=age * 365)).strftime("%Y-%m-%d")

    risk = random.choices(risk_levels, weights=[2, 5, 2, 1])[0]
    patient_type = random.choice(patient_types)
    # assess_result = "A3" if risk == "高危" else ("A4" if risk == "低危" else ("A5" if risk == "无法评估" else "A2"))
    assess_result = random.choice(risk_levels)
    visit_datetime = fake.date_time_between(start_date='-2y', end_date='now')
    discharged_datetime = visit_datetime + timedelta(days=random.randint(1, 10))
    assess_time = visit_datetime + timedelta(days=random.randint(0, 1))
    create_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    return {
        "id": i + 1,
        "project_id": "PRJ001",
        "patient_id": f"PT{i + 1:04d}",
        "visit_sn": f"VSN{i + 1:04d}",
        "cancer_name": '肺癌',
        "hospital_code": "HSP001",
        "hospital_name": "示例医院",
        "department_code": "DEP001",
        "department_name": "肿瘤科",
        "discharged_datetime": discharged_datetime.strftime("%Y-%m-%d %H:%M:%S"),
        "visit_datetime": visit_datetime.strftime("%Y-%m-%d %H:%M:%S"),
        "admission": visit_datetime.strftime("%Y-%m-%d %H:%M:%S"),
        "time": random.randint(1, 10),
        "have_admission_record": random.choice([0, 1]),
        "past_admission_count": random.randint(0, 5),
        "past_outpatient_count": random.randint(0, 10),
        "sex": sex,
        "age": age,
        "birthday": birthday,
        "assess_result": assess_result,
        "assess_time": assess_time.strftime("%Y-%m-%d %H:%M:%S"),
        "high_risks": risk,
        "patient_type_code": patient_type,
        "patient_type": patient_type,
        "past_not_visit_sn": random.choice([0, 1]),
        "past_outpatient_only_have": random.choice([0, 1]),
        "past_admission_only_have": random.choice([0, 1]),
        "past_outpatient_have_adm_have": random.choice([0, 1]),
        "assess_0_6_month": random.choice([0, 1]),
        "assess_12_month": random.choice([0, 1]),
        "create_time": create_time
    }


# 生成 10 条示例
patients = [generate_patient(i) for i in range(10)]

for p in patients:
    print(p)
